package blockworkpool

import (
	"context"
	"sync"
)

// 多coroutine pool 管理
type WorkPoolMgr struct {
	poolMap map[string]*Pool
	wg      sync.WaitGroup
	ctx     context.Context
}

func (mgr *WorkPoolMgr) GoWaitTask(task Task) {
	mgr.GoWaitRunTask(task)
}

func (mgr *WorkPoolMgr) GoWaitRunTask(task RunTask) {
	mgr.wg.Add(1)
	go func() {
		defer mgr.wg.Done()
		task.Run()
	}()
}

func (mgr *WorkPoolMgr) waitProcessTask(task Task) Task {
	mgr.wg.Add(1)
	return func() {
		defer mgr.wg.Done()
		task()
	}
}

// 任务放入执行队列
func (mgr *WorkPoolMgr) Enqueue(namePool string, task Task) {
	mgr.RunTaskEnqueue(namePool, task)
}

// 任务放入执行队列
func (mgr *WorkPoolMgr) RunTaskEnqueue(namePool string, task RunTask) {
	pool, ok := mgr.poolMap[namePool]
	if !ok {
		return
	}
	pool.EnqueueTask(mgr.waitProcessTask(task.Run))
}

// 创建coroutine pool 管理
func NewWorkPoolMgr(ctx context.Context) *WorkPoolMgr {
	return &WorkPoolMgr{
		poolMap: map[string]*Pool{},
		ctx:     ctx,
	}
}

// 添加新的协程池
func (mgr *WorkPoolMgr) AddPool(name string, limit int, bfSize ...int) {
	_, ok := mgr.poolMap[name]
	if ok {
		return
	}
	mgr.poolMap[name] = NewPool(mgr.ctx, limit, bfSize...)
}

// 等待所有任务完成，
// 另外这里在等待的过程中创建新的任务放入
func (mgr *WorkPoolMgr) Wait() {
	mgr.wg.Wait()
}
